import os
import sys
import pickle
import random
import argparse
from tqdm import tqdm
from shutil import copy, rmtree
from torchvision import transforms, datasets
import cv2
import PIL.Image as Image
from matplotlib import pyplot as plt
import random
import shutil
import glob
from loguru import logger
from datetime import datetime
logger.add(f"../log/{str(datetime.today())[:10]}_preprocess_classification_data.log", level='INFO', encoding="utf-8", enqueue=True)


def parse_option():
    parser = argparse.ArgumentParser('preprocess data')
    parser.add_argument('--root_path', 
            type=str, 
            default='../datasets/flowers/',
            help='Modify according to the data name')
    parser.add_argument('--split_ratio',
            type=float,
            default=0.2,
            help='The ratio of trianing and test sets')
    parser.add_argument('--bool_enhance',
            type=bool,
            default=False,
            help='Whether to enrich the dataset')
    parser.add_argument('--thr_enhance',
            type=int,
            default=1300,
            help='The threshold for the number of data enrichments')
    parser.add_argument('--mul_enhance',
            type=float,
            default=1.0,
            help='The multiple for the number of data enrichments')

    args = parser.parse_args()
    return args


# 创建二级目录文件
def mkfile(file_name: str, cls:list):
    if not os.path.exists(file_name):
        os.mkdir(file_name)
        for c in cls:
            if not os.path.exists(os.path.join(file_name, c)):
                print(os.path.join(file_name, c))
                os.mkdir(os.path.join(file_name, c))
                logger.info(f"mkdir {os.path.join(file_name, c)} success")


def split_data(root_path: str, split_ratio:float):
    raw_data = os.path.join(root_path, 'flower')
    train_data = os.path.join(root_path, 'train')
    test_data = os.path.join(root_path, 'test')

    cls = [c for c in os.listdir(raw_data) if '.txt' not in c]
    mkfile(train_data, cls)
    mkfile(test_data, cls)
    
    # 切分
    for c in cls:
        c_path = os.path.join(raw_data, c)
        #imgs = os.listdir(c_path)
        path_imgs = glob.glob(raw_data + '/' + c +'/*')
        num = len(path_imgs)
        index_list = list(range(num))
        random.shuffle(index_list)
        test_index = index_list[:int(num*split_ratio)]
        train_index = index_list[int(num*split_ratio):]
        for i in test_index:
            shutil.copy(path_imgs[i], test_data + '/' + c  + '/')
        for j in train_index:
            shutil.copy(path_imgs[j], train_data + '/' + c  + '/')


def exchange_pic(pic_path):
    img = Image.open(pic_path)
    num = random.randint(0,11)  # 首尾都可能出现
    if num == 0 or num ==1:
        new_img = transforms.CenterCrop(900)(img)
    elif num == 2 or num == 3:
        new_img = transforms.RandomHorizontalFlip(p=1)(img)
    elif num == 4  or num ==5:
        new_img = transforms.RandomHorizontalFlip(p=1)(img)
    elif num == 6:
        new_img = transforms.ColorJitter(brightness=1)(img)
    elif num == 7:
        new_img = transforms.ColorJitter(contrast=2)(img)
    elif num == 8:
        new_img = transforms.ColorJitter(saturation=3)(img)
    elif num == 9:
        new_img = transforms.ColorJitter(hue=0.4)(img)
    else:
        new_img = transforms.Grayscale(num_output_channels=3)(img)
    return  new_img


def enhance_data(root_path, threshold, mul_enhance):  # train
    path = os.path.join(root_path, 'train')
    assert os.path.exists(path), f'{path} does not exist.'
    path_folder = glob.glob(path+'/*')
    nums_max = max([len(glob.glob(i+'/*')) for i in path_folder])
    ex_nums_max = max(int(nums_max*mul_enhance), threshold)
    print(threshold, ex_nums_max, path_folder )
    for i in range(len(path_folder)):
        path_pics = glob.glob(path_folder[i]+'/*.jpg') 
        if len(path_pics) > ex_nums_max:
            continue
        else:
            for j in range(len(glob.glob(path_folder[i]+'/*')), ex_nums_max-1):
                path_one_pic = random.choice(path_pics)
                path_new_pic = ''.join([path_one_pic[:-4], '_', str(j), ".jpg"])
                new_pic = exchange_pic(path_one_pic)
                new_pic.save(path_new_pic)
        logger.info(f"expand {path_folder[i]} success")
        

def main(args):
    root_path = args.root_path
    split_ratio = args.split_ratio
    bool_enhance = args.bool_enhance
    thr_enhance = args.thr_enhance
    mul_enhance = args.mul_enhance

    split_data(root_path, split_ratio)
    if bool_enhance:
        enhance_data(root_path, thr_enhance, mul_enhance)


if __name__ == '__main__':
    args_ = parse_option()
    main(args_)